{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SAC的简化版实现,alpha使用常量代替.只使用一个value模型,而不是两个."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAUeUlEQVR4nO3dfWxT570H8K/txE5CcpwmNHbT2AKp3EHEWxcgOe2VNlGXrI1QGdHVNqE2rRAVzEGl2eVukVr6spdU7I9u3Wj4Z4VqEmM3ldhEREvTQIN2cUhJm6uQlLS9F24ywHYh18dJSuzE57l/cHOGQ2gxSfzE8fcjHQk/zxP7d07iL+ecx8fHJIQQICKSwCy7ACJKXwwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKSRloA7du3D4sWLUJWVhbKy8vR0dEhqxQikkRKAP35z39GXV0dXnrpJXz88cdYtWoVKisrEQwGZZRDRJKYZFyMWl5ejrVr1+L3v/89AEDXdbhcLuzcuRM/+9nPkl0OEUmSkewXjEaj6OzsRH19vdFmNpvh8Xjg8/mm/JlIJIJIJGI81nUdg4ODKCwshMlkmvWaiSgxQggMDQ2huLgYZvPtD7SSHkBXr15FLBaDw+GIa3c4HDh//vyUP9PQ0IBXXnklGeUR0QwaGBhASUnJbfuTHkB3o76+HnV1dcZjTdPgdrsxMDAARVEkVkZEUwmHw3C5XMjLy/vacUkPoIULF8JisSAQCMS1BwIBOJ3OKX/GZrPBZrPd0q4oCgOIaA77plMkSZ8Fs1qtKCsrQ2trq9Gm6zpaW1uhqmqyyyEiiaQcgtXV1aGmpgZr1qzBunXr8Jvf/AYjIyN45plnZJRDRJJICaAf/OAH+PLLL7Fnzx74/X6sXr0a77333i0npolofpPyOaDpCofDsNvt0DSN54CI5qA7fY/yWjAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJk3AAnTp1Chs3bkRxcTFMJhP+8pe/xPULIbBnzx7cd999yM7Ohsfjweeffx43ZnBwEFu2bIGiKMjPz8fWrVsxPDw8rRUhotSTcACNjIxg1apV2Ldv35T9e/fuxRtvvIH9+/fjzJkzWLBgASorKzE6OmqM2bJlC3p6etDS0oLm5macOnUKzz777N2vBRGlJjENAMSRI0eMx7quC6fTKX79618bbaFQSNhsNvGnP/1JCCFEb2+vACA++ugjY8y7774rTCaTuHTp0h29rqZpAoDQNG065RPRLLnT9+iMngO6cOEC/H4/PB6P0Wa321FeXg6fzwcA8Pl8yM/Px5o1a4wxHo8HZrMZZ86cmfJ5I5EIwuFw3EJEqW9GA8jv9wMAHA5HXLvD4TD6/H4/ioqK4vozMjJQUFBgjJmsoaEBdrvdWFwu10yWTUSSpMQsWH19PTRNM5aBgQHZJRHRDJjRAHI6nQCAQCAQ1x4IBIw+p9OJYDAY1z8+Po7BwUFjzGQ2mw2KosQtRJT6ZjSAFi9eDKfTidbWVqMtHA7jzJkzUFUVAKCqKkKhEDo7O40xJ06cgK7rKC8vn8lyiGiOy0j0B4aHh/HFF18Yjy9cuICuri4UFBTA7XZj165d+MUvfoElS5Zg8eLFePHFF1FcXIxNmzYBAJYtW4bvfe972LZtG/bv34+xsTHU1tbihz/8IYqLi2dsxYgoBSQ6vXby5EkB4JalpqZGCHFjKv7FF18UDodD2Gw28cgjj4i+vr6457h27Zr40Y9+JHJzc4WiKOKZZ54RQ0NDMz7FR0Ry3Ol71CSEEBLz766Ew2HY7XZomsbzQURz0J2+R1NiFoyI5icGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCRNwrflIZppQghcH7yE64OXjLaM7Dwo9y+DyWSSWBnNNgYQzQnXPj8D/38eNx7nOpdAuX+ZxIooGXgIRnOAgD4WkV0EScAAIvkEEBtnAKUjBhBJJ4SOWOSruDazhWcH0gEDiKQTQsf10JW4tpxCl6RqKJkYQCSfEDeWm5itWZKKoWRiANGcZMlkAKWDhAKooaEBa9euRV5eHoqKirBp0yb09fXFjRkdHYXX60VhYSFyc3NRXV2NQCAQN6a/vx9VVVXIyclBUVERdu/ejfHx8emvDc0b5kyb7BIoCRIKoLa2Nni9XrS3t6OlpQVjY2PYsGEDRkZGjDHPP/88jh49iqamJrS1teHy5cvYvHmz0R+LxVBVVYVoNIrTp0/j7bffxsGDB7Fnz56ZWytKKULXISYfgpktkqqhpBLTEAwGBQDR1tYmhBAiFAqJzMxM0dTUZIz59NNPBQDh8/mEEEIcO3ZMmM1m4ff7jTGNjY1CURQRiUTu6HU1TRMAhKZp0ymf5ojoiCY++eO/iY7924zly/P/IbssmoY7fY9O6xyQpmkAgIKCAgBAZ2cnxsbG4PF4jDFLly6F2+2Gz+cDAPh8PqxYsQIOh8MYU1lZiXA4jJ6enilfJxKJIBwOxy00f+ixMUCPyS6DJLjrANJ1Hbt27cLDDz+M5cuXAwD8fj+sVivy8/PjxjocDvj9fmPMzeEz0T/RN5WGhgbY7XZjcbk4RTuf6ONR6JMDiNeApYW7DiCv14tz587h8OHDM1nPlOrr66FpmrEMDAzM+mtS8kSHB+M+iGjOsCHLXiSxIkqWu/q4aW1tLZqbm3Hq1CmUlJQY7U6nE9FoFKFQKG4vKBAIwOl0GmM6Ojrinm9ilmxizGQ2mw02G2dF5ishdAD/OAltMpthzsyWVxAlTUJ7QEII1NbW4siRIzhx4gQWL14c119WVobMzEy0trYabX19fejv74eqqgAAVVXR3d2NYDBojGlpaYGiKCgtLZ3OutB8YTLDnJEpuwpKgoT2gLxeLw4dOoS//vWvyMvLM87Z2O12ZGdnw263Y+vWrairq0NBQQEURcHOnTuhqioqKioAABs2bEBpaSmefPJJ7N27F36/Hy+88AK8Xi/3cggAYDKZYc6wyi6DkiChAGpsbAQAfPe7341rP3DgAJ5++mkAwOuvvw6z2Yzq6mpEIhFUVlbizTffNMZaLBY0Nzdjx44dUFUVCxYsQE1NDV599dXprQmlLBGb9CFU040QovnPJMSkT4ClgHA4DLvdDk3ToCiK7HJomq5+1o4LJ98yHmdk52H5v7yMzOw8iVXRdNzpe5T/zZB0Or8LKG0xgEi6WHR0UouJnwNKEwwgku7mL6MHgCzlXpgtnAVLBwwgkk4fj8Y9tmQtgIkXo6YFBhDNOeYMK2/HkyYYQCTVVJOwlgwrwGn4tMDfMkkmICZdiMrDr/TBACKphK7f+DqOSXgIlh4YQCSXmDqAKD0wgEgqXY/dcldUk5n3BEsXDCCSKha9jkj4y7i27IL7JVVDycYAIrmEmDQTZoLFyu8CShcMIJpzLLwlT9pgANGcw3uCpQ8GEEl14+tYb2ICzBaehE4XDCCSSh+P3nJfeEofDCCSSh9jAKUzBhBJFRsbhcDNh2Gm/18oHTCASKrR/70c953Qmdl5yMyxS6yIkokBRFJNvhDVbMnkHTHSCAOI5hSTJYNXw6cRBhDNKSazhdPwaYQBRNIIISD0+M8BmcwW7gGlEQYQSRXjLXnSGvd1aVaNj49jeHh46k4hcH1kKK4pFoshFArddi/IarUiJydnpsskSRhANKu6urqwefNm6JMOtYAbt/7a/vgKPPqgy2jr7e3Fxn9djZg+9YcTn3rqKfzqV7+atXopuRhANKui0SguXbo0ZQBZzCY47GtwObIEf4/8E/Isg/ji4r9j4O9/v+2Ho0Oh0OwWTEmV0DmgxsZGrFy5EoqiQFEUqKqKd9991+gfHR2F1+tFYWEhcnNzUV1djUAgEPcc/f39qKqqQk5ODoqKirB7926Mj49PfilKAyaTCUOmB3Bu+J8xOHY//md0OS6b1oOnJtNHQr/pkpISvPbaa+js7MTZs2exfv16PPHEE+jp6QEAPP/88zh69CiamprQ1taGy5cvY/PmzcbPx2IxVFVVIRqN4vTp03j77bdx8OBB7NmzZ2bXilLGV7E86Ji4C6oJmTn38Qvp00hCh2AbN26Me/zLX/4SjY2NaG9vR0lJCf7whz/g0KFDWL9+PQDgwIEDWLZsGdrb21FRUYH3338fvb29+OCDD+BwOLB69Wr8/Oc/x09/+lO8/PLLsFr5Cdh0U5h5CcNmDV/pCswYR7D/1JSHazQ/3fU5oFgshqamJoyMjEBVVXR2dmJsbAwej8cYs3TpUrjdbvh8PlRUVMDn82HFihVwOBzGmMrKSuzYsQM9PT148MEHE6rh/PnzyM3NvdtVoCS4ePHilDcfBICYruPYqTbYcvrwZdSFBRYN//1fHQBuf3X84OAgent7Z6lamim3nfmcJOEA6u7uhqqqGB0dRW5uLo4cOYLS0lJ0dXXBarUiPz8/brzD4YDf7wcA+P3+uPCZ6J/ou51IJIJI5B+fFwmHwwAATdN4/miO+7o/RCGAo6c/A/DZHT9fNBrliegUMDIyckfjEg6gb33rW+jq6oKmaXjnnXdQU1ODtra2hAtMRENDA1555ZVb2svLy6Eoyqy+Nk2fyWS67V5QopxOJx566KEZeS6aPRM7Cd8k4ekGq9WKBx54AGVlZWhoaMCqVavw29/+Fk6nc8r/nQKBAJxOJ4AbfzyTZ8UmHk+MmUp9fT00TTOWgYGBRMsmojlo2vOduq4jEomgrKwMmZmZaG1tNfr6+vrQ398PVVUBAKqqoru7G8Fg0BjT0tICRVFQWlp629ew2WzG1P/EQkSpL6FDsPr6ejz22GNwu90YGhrCoUOH8OGHH+L48eOw2+3YunUr6urqUFBQAEVRsHPnTqiqioqKCgDAhg0bUFpaiieffBJ79+6F3+/HCy+8AK/XC5uNd0IgSjcJBVAwGMRTTz2FK1euwG63Y+XKlTh+/DgeffRRAMDrr78Os9mM6upqRCIRVFZW4s033zR+3mKxoLm5GTt27ICqqliwYAFqamrw6quvzuxa0ZxhsVigKMqMTa1nZWXNyPPQ3GASM3V2MInC4TDsdjs0TePh2BwXiUTiDrmnKzc3F/fcc8+MPR/Njjt9j/JaMJpVNpsNLpfrmwdSWuJFN0QkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEiaDNkF3A0hBAAgHA5LroSIpjLx3px4r95OSgbQtWvXAAAul0tyJUT0dYaGhmC322/bn5IBVFBQAADo7+//2pWjeOFwGC6XCwMDA1AURXY5KYHb7O4IITA0NITi4uKvHZeSAWQ23zh1Zbfb+UdxFxRF4XZLELdZ4u5k54AnoYlIGgYQEUmTkgFks9nw0ksvwWazyS4lpXC7JY7bbHaZxDfNkxERzZKU3AMiovmBAURE0jCAiEgaBhARSZOSAbRv3z4sWrQIWVlZKC8vR0dHh+ySpGloaMDatWuRl5eHoqIibNq0CX19fXFjRkdH4fV6UVhYiNzcXFRXVyMQCMSN6e/vR1VVFXJyclBUVITdu3djfHw8masizWuvvQaTyYRdu3YZbdxmSSJSzOHDh4XVahVvvfWW6OnpEdu2bRP5+fkiEAjILk2KyspKceDAAXHu3DnR1dUlHn/8ceF2u8Xw8LAxZvv27cLlconW1lZx9uxZUVFRIR566CGjf3x8XCxfvlx4PB7xySefiGPHjomFCxeK+vp6GauUVB0dHWLRokVi5cqV4rnnnjPauc2SI+UCaN26dcLr9RqPY7GYKC4uFg0NDRKrmjuCwaAAINra2oQQQoRCIZGZmSmampqMMZ9++qkAIHw+nxBCiGPHjgmz2Sz8fr8xprGxUSiKIiKRSHJXIImGhobEkiVLREtLi/jOd75jBBC3WfKk1CFYNBpFZ2cnPB6P0WY2m+HxeODz+SRWNndomgbgHxfsdnZ2YmxsLG6bLV26FG6329hmPp8PK1asgMPhMMZUVlYiHA6jp6cnidUnl9frRVVVVdy2AbjNkimlLka9evUqYrFY3C8dABwOB86fPy+pqrlD13Xs2rULDz/8MJYvXw4A8Pv9sFqtyM/PjxvrcDjg9/uNMVNt04m++ejw4cP4+OOP8dFHH93Sx22WPCkVQPT1vF4vzp07h7/97W+yS5nTBgYG8Nxzz6GlpQVZWVmyy0lrKXUItnDhQlgslltmIwKBAJxOp6Sq5oba2lo0Nzfj5MmTKCkpMdqdTiei0ShCoVDc+Ju3mdPpnHKbTvTNN52dnQgGg/j2t7+NjIwMZGRkoK2tDW+88QYyMjLgcDi4zZIkpQLIarWirKwMra2tRpuu62htbYWqqhIrk0cIgdraWhw5cgQnTpzA4sWL4/rLysqQmZkZt836+vrQ399vbDNVVdHd3Y1gMGiMaWlpgaIoKC0tTc6KJNEjjzyC7u5udHV1GcuaNWuwZcsW49/cZkki+yx4og4fPixsNps4ePCg6O3tFc8++6zIz8+Pm41IJzt27BB2u118+OGH4sqVK8by1VdfGWO2b98u3G63OHHihDh79qxQVVWoqmr0T0wpb9iwQXR1dYn33ntP3HvvvWk1pXzzLJgQ3GbJknIBJIQQv/vd74Tb7RZWq1WsW7dOtLe3yy5JGgBTLgcOHDDGXL9+Xfz4xz8W99xzj8jJyRHf//73xZUrV+Ke5+LFi+Kxxx4T2dnZYuHCheInP/mJGBsbS/LayDM5gLjNkoNfx0FE0qTUOSAiml8YQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTT/ByUzzqKpfR1YAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "\n",
    "    def __init__(self):\n",
    "        env = gym.make('CartPole-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        over = terminated or truncated\n",
    "\n",
    "        #限制最大步数\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            over = True\n",
    "\n",
    "        return state, reward, over\n",
    "\n",
    "    #打印游戏图像\n",
    "    def show(self):\n",
    "        from matplotlib import pyplot as plt\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.imshow(self.env.render())\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()\n",
    "\n",
    "env.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/cuda117/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[0.5274, 0.4726],\n",
       "        [0.6016, 0.3984]], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "model_action = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    "    torch.nn.Softmax(dim=1),\n",
    ")\n",
    "\n",
    "model_action(torch.randn(2, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.1879, -0.0926],\n",
       "        [-0.2231, -0.0472]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_value = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "model_value_next = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "model_value_next.load_state_dict(model_value.state_dict())\n",
    "\n",
    "model_value(torch.randn(2, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/cuda117/lib/python3.10/site-packages/gym/utils/passive_env_checker.py:233: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  if not isinstance(terminated, (bool, np.bool8)):\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "18.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "import random\n",
    "\n",
    "\n",
    "#玩一局游戏并记录数据\n",
    "def play(show=False):\n",
    "    data = []\n",
    "    reward_sum = 0\n",
    "\n",
    "    state = env.reset()\n",
    "    over = False\n",
    "    while not over:\n",
    "        prob = model_action(torch.FloatTensor(state).reshape(1, 4))[0].tolist()\n",
    "        action = random.choices(range(2), weights=prob, k=1)[0]\n",
    "\n",
    "        next_state, reward, over = env.step(action)\n",
    "\n",
    "        data.append((state, action, reward, next_state, over))\n",
    "        reward_sum += reward\n",
    "\n",
    "        state = next_state\n",
    "\n",
    "        if show:\n",
    "            display.clear_output(wait=True)\n",
    "            env.show()\n",
    "\n",
    "    return data, reward_sum\n",
    "\n",
    "\n",
    "play()[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_20457/3891364554.py:27: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ../torch/csrc/utils/tensor_new.cpp:230.)\n",
      "  state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(210,\n",
       " (array([0.02908172, 0.02583727, 0.01217754, 0.02771634], dtype=float32),\n",
       "  0,\n",
       "  1.0,\n",
       "  array([ 0.02959846, -0.16945718,  0.01273187,  0.3242164 ], dtype=float32),\n",
       "  False))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#数据池\n",
    "class Pool:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.pool = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.pool)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.pool[i]\n",
    "\n",
    "    #更新动作池\n",
    "    def update(self):\n",
    "        #每次更新不少于N条新数据\n",
    "        old_len = len(self.pool)\n",
    "        while len(pool) - old_len < 200:\n",
    "            self.pool.extend(play()[0])\n",
    "\n",
    "        #只保留最新的N条数据\n",
    "        self.pool = self.pool[-2_0000:]\n",
    "\n",
    "    #获取一批数据样本\n",
    "    def sample(self):\n",
    "        data = random.sample(self.pool, 64)\n",
    "\n",
    "        state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n",
    "        action = torch.LongTensor([i[1] for i in data]).reshape(-1, 1)\n",
    "        reward = torch.FloatTensor([i[2] for i in data]).reshape(-1, 1)\n",
    "        next_state = torch.FloatTensor([i[3] for i in data]).reshape(-1, 4)\n",
    "        over = torch.LongTensor([i[4] for i in data]).reshape(-1, 1)\n",
    "\n",
    "        return state, action, reward, next_state, over\n",
    "\n",
    "\n",
    "pool = Pool()\n",
    "pool.update()\n",
    "state, action, reward, next_state, over = pool.sample()\n",
    "\n",
    "len(pool), pool[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_action = torch.optim.Adam(model_action.parameters(), lr=2e-4)\n",
    "optimizer_value = torch.optim.Adam(model_value.parameters(), lr=2e-3)\n",
    "\n",
    "\n",
    "def soft_update(_from, _to):\n",
    "    for _from, _to in zip(_from.parameters(), _to.parameters()):\n",
    "        value = _to.data * 0.995 + _from.data * 0.005\n",
    "        _to.data.copy_(value)\n",
    "\n",
    "\n",
    "def get_prob_entropy(state):\n",
    "    prob = model_action(torch.FloatTensor(state).reshape(-1, 4))\n",
    "    entropy = prob * (prob + 1e-8).log()\n",
    "    entropy = -entropy.sum(dim=1, keepdim=True)\n",
    "\n",
    "    return prob, entropy\n",
    "\n",
    "\n",
    "def requires_grad(model, value):\n",
    "    for param in model.parameters():\n",
    "        param.requires_grad_(value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0156246423721313"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def train_value(state, action, reward, next_state, over):\n",
    "    requires_grad(model_value, True)\n",
    "    requires_grad(model_action, False)\n",
    "\n",
    "    #计算target\n",
    "    with torch.no_grad():\n",
    "        #计算动作的熵\n",
    "        prob, entropy = get_prob_entropy(next_state)\n",
    "        target = model_value_next(next_state)\n",
    "\n",
    "    #加权熵,熵越大越好\n",
    "    target = (prob * target).sum(dim=1, keepdim=True)\n",
    "    target = target + 1e-3 * entropy\n",
    "    target = target * 0.98 * (1 - over) + reward\n",
    "\n",
    "    #计算value\n",
    "    value = model_value(state).gather(dim=1, index=action)\n",
    "    loss = torch.nn.functional.mse_loss(value, target)\n",
    "    loss.backward()\n",
    "    optimizer_value.step()\n",
    "    optimizer_value.zero_grad()\n",
    "\n",
    "    return loss.item()\n",
    "\n",
    "\n",
    "train_value(state, action, reward, next_state, over)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.09280812740325928"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def train_action(state):\n",
    "    requires_grad(model_value, False)\n",
    "    requires_grad(model_action, True)\n",
    "\n",
    "    #计算熵\n",
    "    prob, entropy = get_prob_entropy(state)\n",
    "\n",
    "    #计算value\n",
    "    value = model_value(state)\n",
    "\n",
    "    #求期望求和\n",
    "    value = (prob * value).sum(dim=1, keepdim=True)\n",
    "\n",
    "    #加权熵\n",
    "    loss = -(value + 1e-3 * entropy).mean()\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer_action.step()\n",
    "    optimizer_action.zero_grad()\n",
    "\n",
    "    return loss.item()\n",
    "\n",
    "\n",
    "train_action(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "OHoSU6uI-xIt",
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 429 14.5\n",
      "10 2735 183.65\n",
      "20 5004 198.0\n",
      "30 7707 180.8\n",
      "40 11235 182.7\n",
      "50 14310 179.65\n",
      "60 17349 183.85\n",
      "70 20000 184.35\n",
      "80 20000 187.75\n",
      "90 20000 155.85\n",
      "100 20000 130.05\n",
      "110 20000 161.45\n",
      "120 20000 154.8\n",
      "130 20000 192.6\n",
      "140 20000 195.2\n",
      "150 20000 149.15\n",
      "160 20000 170.0\n",
      "170 20000 193.35\n",
      "180 20000 199.0\n",
      "190 20000 200.0\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    model_action.train()\n",
    "    model_value.train()\n",
    "\n",
    "    #训练N次\n",
    "    for epoch in range(200):\n",
    "        #更新N条数据\n",
    "        pool.update()\n",
    "\n",
    "        #每次更新过数据后,学习N次\n",
    "        for i in range(200):\n",
    "            #采样一批数据\n",
    "            state, action, reward, next_state, over = pool.sample()\n",
    "\n",
    "            #训练\n",
    "            train_value(state, action, reward, next_state, over)\n",
    "            train_action(state)\n",
    "            soft_update(model_value, model_value_next)\n",
    "\n",
    "        if epoch % 10 == 0:\n",
    "            test_result = sum([play()[-1] for _ in range(20)]) / 20\n",
    "            print(epoch, len(pool), test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAT0klEQVR4nO3dfWxT97kH8K+dF4e8HKchjb0o8QWtvWMRL9sCJGeVbtfVJWPRNNZcaZu4XVYhUJmDSlNxtUiUCtSrcNmV2nWj4Y9pwP7omDKJTs1oqyi0QROGQNpoIYVo3bolK9jmRTl2QmMn9nP/6HJaEwN24vgX0+9HOlJ9fo/t55xyvpzzOza2iIiAiEgBq+oGiOjziwFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKKAuggwcPYtmyZSgoKEBdXR36+vpUtUJEiigJoN/97ndobW3F888/j3fffRdr1qxBQ0MDAoGAinaISBGLii+j1tXVYd26dfjlL38JAIjFYqiursaOHTvw05/+NNPtEJEiuZl+w0gkgv7+frS1tZnrrFYr3G43vF5vwueEw2GEw2HzcSwWw40bN7B06VJYLJYF75mIUiMiCIVCqKyshNV6+wutjAfQtWvXEI1G4XA44tY7HA5cunQp4XPa29uxd+/eTLRHRGk0OjqKqqqq245nPIDmoq2tDa2treZjwzDgcrkwOjoKTdMUdkZEiQSDQVRXV6OkpOSOdRkPoPLycuTk5MDv98et9/v9cDqdCZ9js9lgs9lmrdc0jQFEtIjdbYok43fB8vPzUVtbi56eHnNdLBZDT08PdF3PdDtEpJCSS7DW1lY0Nzdj7dq1WL9+PV566SVMTEzgySefVNEOESmiJIC+//3v4+rVq9izZw98Ph++8pWv4M0335w1MU1E9zYlnwOar2AwCLvdDsMwOAdEtAgle4zyu2BEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJSJuUAOnXqFL7zne+gsrISFosFr732Wty4iGDPnj34whe+gCVLlsDtduMvf/lLXM2NGzewefNmaJqG0tJSbNmyBePj4/PaECLKPikH0MTEBNasWYODBw8mHD9w4ABefvllHDp0CGfPnkVRUREaGhowOTlp1mzevBlDQ0Po7u5GV1cXTp06hW3bts19K4goO8k8AJDjx4+bj2OxmDidTvnZz35mrhsbGxObzSa//e1vRUTk/fffFwBy7tw5s+aNN94Qi8UiH330UVLvaxiGABDDMObTPhEtkGSP0bTOAX344Yfw+Xxwu93mOrvdjrq6Oni9XgCA1+tFaWkp1q5da9a43W5YrVacPXs24euGw2EEg8G4hYiyX1oDyOfzAQAcDkfceofDYY75fD5UVFTEjefm5qKsrMysuVV7ezvsdru5VFdXp7NtIlIkK+6CtbW1wTAMcxkdHVXdEhGlQVoDyOl0AgD8fn/cer/fb445nU4EAoG48enpady4ccOsuZXNZoOmaXELEWW/tAbQ8uXL4XQ60dPTY64LBoM4e/YsdF0HAOi6jrGxMfT395s1J0+eRCwWQ11dXTrbIaJFLjfVJ4yPj+ODDz4wH3/44YcYGBhAWVkZXC4Xdu7ciRdeeAEPPvggli9fjueeew6VlZXYtGkTAODLX/4yvvWtb2Hr1q04dOgQpqam0NLSgh/84AeorKxM24YRURZI9fba22+/LQBmLc3NzSLyya345557ThwOh9hsNnn00UdleHg47jWuX78uP/zhD6W4uFg0TZMnn3xSQqFQ2m/xEZEayR6jFhERhfk3J8FgEHa7HYZhcD6IaBFK9hjNirtgRHRvYgARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyqT8szxEdxOLRWH8YxCx6bC5rtj5AGwlSxV2RYsRA4jSLjYdwcjpY4iM3zDXfdG9jQFEs/ASjNJP5JOF6C4YQJR+ImD8UDIYQLQAGD+UHAYQpd0nP7bLEKK7YwDRwmD+UBIYQJR+PAOiJKUUQO3t7Vi3bh1KSkpQUVGBTZs2YXh4OK5mcnISHo8HS5cuRXFxMZqamuD3++NqRkZG0NjYiMLCQlRUVGDXrl2Ynp6e/9bQosApaEpWSgHU29sLj8eDM2fOoLu7G1NTU9iwYQMmJibMmmeeeQavv/46Ojs70dvbi8uXL+Pxxx83x6PRKBobGxGJRHD69GkcPXoUR44cwZ49e9K3VaQWb8FTsmQeAoGAAJDe3l4RERkbG5O8vDzp7Ow0ay5evCgAxOv1iojIiRMnxGq1is/nM2s6OjpE0zQJh8NJva9hGAJADMOYT/u0QMKhG/Le0Wel79BWc7n+wTnVbVEGJXuMzmsOyDAMAEBZWRkAoL+/H1NTU3C73WbNihUr4HK54PV6AQBerxerVq2Cw+EwaxoaGhAMBjE0NJTwfcLhMILBYNxCi5fwIoySNOcAisVi2LlzJx566CGsXLkSAODz+ZCfn4/S0tK4WofDAZ/PZ9Z8NnxmxmfGEmlvb4fdbjeX6urqubZNmcBLMErSnAPI4/HgwoULOHbsWDr7SaitrQ2GYZjL6Ojogr8nzRdDiO5uTl9GbWlpQVdXF06dOoWqqipzvdPpRCQSwdjYWNxZkN/vh9PpNGv6+vriXm/mLtlMza1sNhtsNttcWiUVRJg/lJSUzoBEBC0tLTh+/DhOnjyJ5cuXx43X1tYiLy8PPT095rrh4WGMjIxA13UAgK7rGBwcRCAQMGu6u7uhaRpqamrmsy20SHwyA8QEortL6QzI4/Hg1VdfxR/+8AeUlJSYczZ2ux1LliyB3W7Hli1b0NrairKyMmiahh07dkDXddTX1wMANmzYgJqaGjzxxBM4cOAAfD4fdu/eDY/Hw7OcewXngChJKQVQR0cHAOAb3/hG3PrDhw/jxz/+MQDgxRdfhNVqRVNTE8LhMBoaGvDKK6+YtTk5Oejq6sL27duh6zqKiorQ3NyMffv2zW9LaPFgAFGSLCLZ96clGAzCbrfDMAxomqa6HbrFx2M+XHrtfzEd/vQDql90b0PZF9cq7IoyKdljlN8Fo/TjvwdESWIAEZEyDCBKP34bnpLEAKK0E/BzQJQcBhClH8+AKEkMIEo7kdjsW/EWi5pmaFFjAFHaRcZvIDr16Y8SWvNsyC8uU9gRLVYMIEo7icXw2Uswi8UKizVHXUO0aDGAKCMs4CUYzcYAoszgHBAlwACiDLAwgCghBhBlBC/BKBEGEC08C3gGRAkxgCgDLP9aiOIxgCgjeAJEiTCAaMFZACYQJcQAooVn4SUYJcYAoszgGRAlwACijOBteEqEAUQZwA8iUmIMIFp4/BwQ3QYDiDLAwkswSogBRJnBMyBKgAFEC47RQ7eT0i+jEgFANBpFKBS67fjNmxNxj2MiCIZCyJ1KXJ+bm4uioiJYeJb0ucMAopT985//xGOPPYabN28mHK9f4cB//2et+fjq1av4r/94GMGbkYT1jzzyCH7zm98sSK+0uDGAKGXRaBSXL1/GxMREwvHr9+fj6lQ1/vHxStisN1E6dQqXr1yGMR5OXH/9+kK2S4tYSnNAHR0dWL16NTRNg6Zp0HUdb7zxhjk+OTkJj8eDpUuXori4GE1NTfD7/XGvMTIygsbGRhQWFqKiogK7du3C9PR0eraGFgVj+n78OfQIrk1V46Pwv+PPoYcRjfHvOpotpQCqqqrC/v370d/fj/Pnz+Ob3/wmvvvd72JoaAgA8Mwzz+D1119HZ2cnent7cfnyZTz++OPm86PRKBobGxGJRHD69GkcPXoUR44cwZ49e9K7VaTUZKwIU1Lwr0cWTEzbERXO71ACMk/33Xef/OpXv5KxsTHJy8uTzs5Oc+zixYsCQLxer4iInDhxQqxWq/h8PrOmo6NDNE2TcDic9HsahiEAxDCM+bZPc/DXv/5VioqKZn59cNZSv3qlvLT/j7L3hT7Z98IZ+b/dL0nRkoLb1m/cuFFisZjqzaI0SvYYnfN5cTQaRWdnJyYmJqDrOvr7+zE1NQW3223WrFixAi6XC16vF/X19fB6vVi1ahUcDodZ09DQgO3bt2NoaAhf/epXU+rh0qVLKC4unusm0ByNjo4iFovddvyyfxTenv+BP/JvsFk+RpF8gMhU4gloABgfH8fFixcXolVSZHx8PKm6lANocHAQuq5jcnISxcXFOH78OGpqajAwMID8/HyUlpbG1TscDvh8PgCAz+eLC5+Z8Zmx2wmHwwiHP53ADAaDAADDMDh/pEAoFILc+sunnzHiNzDSfRrA6aReb3p6GmNjY+lpjhaF292guFXKAfSlL30JAwMDMAwDv//979Hc3Ize3t6UG0xFe3s79u7dO2t9XV0dNE1b0Pem2f72t78hJyd9PzRYWloKXdf5OaB7yMxJwt2k/Eno/Px8PPDAA6itrUV7ezvWrFmDn//853A6nYhEIrP+JvP7/XA6nQAAp9M5667YzOOZmkTa2tpgGIa5jI6Opto2ES1C8/4qRiwWQzgcRm1tLfLy8tDT02OODQ8PY2RkBLquAwB0Xcfg4CACgYBZ093dDU3TUFNTc9v3sNls5q3/mYWIsl9Kl2BtbW3YuHEjXC4XQqEQXn31Vbzzzjt46623YLfbsWXLFrS2tqKsrAyapmHHjh3QdR319fUAgA0bNqCmpgZPPPEEDhw4AJ/Ph927d8Pj8cBmsy3IBhLR4pVSAAUCAfzoRz/ClStXYLfbsXr1arz11lt47LHHAAAvvvgirFYrmpqaEA6H0dDQgFdeecV8fk5ODrq6urB9+3bouo6ioiI0Nzdj37596d0qWlAWiwUlJSVpmwcqLCxMy+tQ9rHInW5nLFLBYBB2ux2GYfByTIHp6Wn4fL473glLRUFBAcrLyzkJfQ9J9hjl5+MpZbm5uaiqqlLdBt0D+O8BEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImVyVTcwFyICAAgGg4o7IaJEZo7NmWP1drIygK5fvw4AqK6uVtwJEd1JKBSC3W6/7XhWBlBZWRkAYGRk5I4bR/GCwSCqq6sxOjoKTdNUt5MVuM/mRkQQCoVQWVl5x7qsDCCr9ZOpK7vdzj8Uc6BpGvdbirjPUpfMyQEnoYlIGQYQESmTlQFks9nw/PPPw2azqW4lq3C/pY77bGFZ5G73yYiIFkhWngER0b2BAUREyjCAiEgZBhARKZOVAXTw4EEsW7YMBQUFqKurQ19fn+qWlGlvb8e6detQUlKCiooKbNq0CcPDw3E1k5OT8Hg8WLp0KYqLi9HU1AS/3x9XMzIygsbGRhQWFqKiogK7du3C9PR0JjdFmf3798NisWDnzp3mOu6zDJEsc+zYMcnPz5df//rXMjQ0JFu3bpXS0lLx+/2qW1OioaFBDh8+LBcuXJCBgQH59re/LS6XS8bHx82ap556Sqqrq6Wnp0fOnz8v9fX18vWvf90cn56elpUrV4rb7Zb33ntPTpw4IeXl5dLW1qZikzKqr69Pli1bJqtXr5ann37aXM99lhlZF0Dr168Xj8djPo5Go1JZWSnt7e0Ku1o8AoGAAJDe3l4RERkbG5O8vDzp7Ow0ay5evCgAxOv1iojIiRMnxGq1is/nM2s6OjpE0zQJh8OZ3YAMCoVC8uCDD0p3d7c8/PDDZgBxn2VOVl2CRSIR9Pf3w+12m+usVivcbje8Xq/CzhYPwzAAfPqF3f7+fkxNTcXtsxUrVsDlcpn7zOv1YtWqVXA4HGZNQ0MDgsEghoaGMth9Znk8HjQ2NsbtG4D7LJOy6suo165dQzQajfufDgAOhwOXLl1S1NXiEYvFsHPnTjz00ENYuXIlAMDn8yE/Px+lpaVxtQ6HAz6fz6xJtE9nxu5Fx44dw7vvvotz587NGuM+y5ysCiC6M4/HgwsXLuBPf/qT6lYWtdHRUTz99NPo7u5GQUGB6nY+17LqEqy8vBw5OTmz7kb4/X44nU5FXS0OLS0t6Orqwttvv42qqipzvdPpRCQSwdjYWFz9Z/eZ0+lMuE9nxu41/f39CAQC+NrXvobc3Fzk5uait7cXL7/8MnJzc+FwOLjPMiSrAig/Px+1tbXo6ekx18ViMfT09EDXdYWdqSMiaGlpwfHjx3Hy5EksX748bry2thZ5eXlx+2x4eBgjIyPmPtN1HYODgwgEAmZNd3c3NE1DTU1NZjYkgx599FEMDg5iYGDAXNauXYvNmzeb/819liGqZ8FTdezYMbHZbHLkyBF5//33Zdu2bVJaWhp3N+LzZPv27WK32+Wdd96RK1eumMvNmzfNmqeeekpcLpecPHlSzp8/L7qui67r5vjMLeUNGzbIwMCAvPnmm3L//fd/rm4pf/YumAj3WaZkXQCJiPziF78Ql8sl+fn5sn79ejlz5ozqlpQBkHA5fPiwWfPxxx/LT37yE7nvvvuksLBQvve978mVK1fiXufvf/+7bNy4UZYsWSLl5eXy7LPPytTUVIa3Rp1bA4j7LDP4z3EQkTJZNQdERPcWBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKTM/wPKvJuX+j9TTwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "play(True)[-1]"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "第7章-DQN算法.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python [conda env:cuda117]",
   "language": "python",
   "name": "conda-env-cuda117-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
